from .AccuracyReward import accuracy_reward
from .FormatReward import format_reward
from .ReasoningStepsReward import reasoning_steps_reward
from .GetCosineScaledReward import get_consine_scaled_reward
from .GetRepetitionPenaltyReward import get_repetition_penalty_reward

def get_reward_functions(script_args):
    reward_funcs_list = []
    reward_funcs_registry = {
        "accuracy": accuracy_reward,
        "format": format_reward,
        "reasoning_steps": reasoning_steps_reward, 
        "consine": get_consine_scaled_reward(
            min_value_wrong=script_args.cosine_min_value_wrong,
            max_value_wrong=script_args.cosine_max_value_wrong,
            min_value_correct=script_args.cosine_min_value_correct,
            max_value_correct=script_args.cosine_max_value_correct,
            max_len=script_args.cosine_max_len,
        ),
        "repetition_penalty": get_repetition_penalty_reward(
            ngram_size=script_args.repetition_n_grams,
            max_penalty=script_args.repetition_max_penalty,
        ),
    }
    for func_name in script_args.reward_funcs:
        if func_name not in reward_funcs_registry:
            raise ValueError(f"Reward function '{func_name}'not found in registry.")
        reward_funcs_list.append(reward_funcs_registry[func_name])
    return reward_funcs_list